# encoding=utf-8

import os

import pandas as pd
from config import config
from core.common import get_data_dict


def dispose_data():
    """处理数据"""
    data_list = get_data_dict()
    for item in data_list:
        if not item.get('is_need'):
            continue
        file_name = f'{item.get("code")}_{item.get("describe")}.xlsx'
        file_path = os.path.join(config.ORIGINAL_PATH, item.get('cate_type'), file_name)
        if not os.path.isfile(file_path):
            continue
        df_data = pd.read_excel(file_path)
        df_data['code'] = item.get("code")
        df_data['describe'] = item.get("describe")
        func = globals()[item.get('indicator_method')]
        func(item, df_data)


def indicator_stock_index(item, pd_data):
    """处理股票指数"""
    pd_data = pd_data.rename(columns={
        '日期': 'Date',
        '开盘': 'Open',
        '收盘': 'Close',
        '最高': 'High',
        '最低': 'Low',
        '成交量': 'Volume'
    })
    cooked_path = os.path.join(config.COOKED_PATH, item.get('cate_type'))
    os.makedirs(cooked_path, exist_ok=True)
    file_name = f'{item.get("code")}_{item.get("describe")}.xlsx'
    pd_data.to_excel(os.path.join(cooked_path, file_name), index=False)


def indicator_etf(item, pd_data):
    """处理ETF"""
    pd_data = pd_data.rename(columns={
        '日期': 'Date',
        '开盘': 'Open',
        '收盘': 'Close',
        '最高': 'High',
        '最低': 'Low',
        '成交量': 'Volume'
    })
    pd_data['净值日期'] = pd.to_datetime(pd_data['净值日期'])
    pd_data['MA1'] = pd_data['累计净值'].rolling(window=10).mean()
    pd_data['MA1_diff'] = pd_data['MA1'] - pd_data['MA1'].shift(1)
    pd_data['MA2'] = pd_data['累计净值'].rolling(window=20).mean()
    pd_data['MA2_diff'] = pd_data['MA2'] - pd_data['MA2'].shift(1)
    pd_data['MA3'] = pd_data['累计净值'].rolling(window=30).mean()
    pd_data['MA3_diff'] = pd_data['MA3'] - pd_data['MA3'].shift(1)

    cooked_path = os.path.join(config.COOKED_PATH, item.get('cate_type'))
    os.makedirs(cooked_path, exist_ok=True)
    weekly_cooked_path = os.path.join(cooked_path, 'weekly')
    os.makedirs(weekly_cooked_path, exist_ok=True)
    file_name = f'{item.get("code")}_{item.get("describe")}.xlsx'
    pd_data.to_excel(os.path.join(weekly_cooked_path, file_name), index=False)


def indicator_active_bond(item, pd_data):
    """处理积极债"""
    cooked_path = os.path.join(config.COOKED_PATH, item.get('cate_type'))
    os.makedirs(cooked_path, exist_ok=True)
    weekly_cooked_path = os.path.join(cooked_path, 'weekly')
    os.makedirs(weekly_cooked_path, exist_ok=True)
    daily_cooked_path = os.path.join(cooked_path, 'daily')
    os.makedirs(daily_cooked_path, exist_ok=True)
    file_name = f'{item.get("code")}_{item.get("describe")}.xlsx'
    weekly_pd_data = calculate_weekly_chat(pd_data)
    weekly_pd_data = get_draw_down(weekly_pd_data)
    weekly_pd_data['MA1'] = weekly_pd_data['累计净值'].rolling(window=10).mean()
    weekly_pd_data['MA1_diff'] = weekly_pd_data['MA1'] - weekly_pd_data['MA1'].shift(1)
    weekly_pd_data['MA2'] = weekly_pd_data['累计净值'].rolling(window=20).mean()
    weekly_pd_data['MA2_diff'] = weekly_pd_data['MA2'] - weekly_pd_data['MA2'].shift(1)
    weekly_pd_data['MA3'] = weekly_pd_data['累计净值'].rolling(window=30).mean()
    weekly_pd_data['MA3_diff'] = weekly_pd_data['MA3'] - weekly_pd_data['MA3'].shift(1)
    pd_data = get_draw_down(pd_data)
    pd_data['MA1'] = pd_data['累计净值'].rolling(window=20).mean()
    pd_data['MA1_diff'] = pd_data['MA1'] - pd_data['MA1'].shift(1)
    pd_data['MA2'] = pd_data['累计净值'].rolling(window=30).mean()
    pd_data['MA2_diff'] = pd_data['MA2'] - pd_data['MA2'].shift(1)
    pd_data['MA3'] = pd_data['累计净值'].rolling(window=40).mean()
    pd_data['MA3_diff'] = pd_data['MA3'] - pd_data['MA3'].shift(1)
    pd_data[40:].to_excel(os.path.join(daily_cooked_path, file_name), index=False)
    weekly_pd_data[40:].to_excel(os.path.join(weekly_cooked_path, file_name), index=False)


def indicator_pure_bond(item, pd_data):
    """处理纯债"""
    cooked_path = os.path.join(config.COOKED_PATH, item.get('cate_type'))
    os.makedirs(cooked_path, exist_ok=True)
    weekly_cooked_path = os.path.join(cooked_path, 'weekly')
    os.makedirs(weekly_cooked_path, exist_ok=True)
    daily_cooked_path = os.path.join(cooked_path, 'daily')
    os.makedirs(daily_cooked_path, exist_ok=True)
    file_name = f'{item.get("code")}_{item.get("describe")}.xlsx'
    weekly_pd_data = calculate_weekly_chat(pd_data)
    weekly_pd_data = get_draw_down(weekly_pd_data)
    weekly_pd_data['MA1'] = weekly_pd_data['累计净值'].rolling(window=5).mean()
    weekly_pd_data['MA1_diff'] = weekly_pd_data['MA1'] - weekly_pd_data['MA1'].shift(1)
    weekly_pd_data['MA2'] = weekly_pd_data['累计净值'].rolling(window=10).mean()
    weekly_pd_data['MA2_diff'] = weekly_pd_data['MA2'] - weekly_pd_data['MA2'].shift(1)
    weekly_pd_data['MA3'] = weekly_pd_data['累计净值'].rolling(window=15).mean()
    weekly_pd_data['MA3_diff'] = weekly_pd_data['MA3'] - weekly_pd_data['MA3'].shift(1)
    pd_data = get_draw_down(pd_data)
    pd_data['MA1'] = pd_data['累计净值'].rolling(window=10).mean()
    pd_data['MA1_diff'] = pd_data['MA1'] - pd_data['MA1'].shift(1)
    pd_data['MA2'] = pd_data['累计净值'].rolling(window=15).mean()
    pd_data['MA2_diff'] = pd_data['MA2'] - pd_data['MA2'].shift(1)
    pd_data['MA3'] = pd_data['累计净值'].rolling(window=20).mean()
    pd_data['MA3_diff'] = pd_data['MA3'] - pd_data['MA3'].shift(1)
    pd_data[20:].to_excel(os.path.join(daily_cooked_path, file_name), index=False)
    weekly_pd_data[20:].to_excel(os.path.join(weekly_cooked_path, file_name), index=False)


def indicator_stock(item, pd_data):
    """处理股票"""
    cooked_path = os.path.join(config.COOKED_PATH, item.get('cate_type'))
    os.makedirs(cooked_path, exist_ok=True)
    weekly_cooked_path = os.path.join(cooked_path, 'weekly')
    os.makedirs(weekly_cooked_path, exist_ok=True)
    daily_cooked_path = os.path.join(cooked_path, 'daily')
    os.makedirs(daily_cooked_path, exist_ok=True)
    file_name = f'{item.get("code")}_{item.get("describe")}.xlsx'
    weekly_pd_data = calculate_weekly_chat(pd_data)
    weekly_pd_data = get_draw_down(weekly_pd_data)
    weekly_pd_data['MA1'] = weekly_pd_data['累计净值'].rolling(window=10).mean()
    weekly_pd_data['MA1_diff'] = weekly_pd_data['MA1'] - weekly_pd_data['MA1'].shift(1)
    weekly_pd_data['MA2'] = weekly_pd_data['累计净值'].rolling(window=20).mean()
    weekly_pd_data['MA2_diff'] = weekly_pd_data['MA2'] - weekly_pd_data['MA2'].shift(1)
    weekly_pd_data['MA3'] = weekly_pd_data['累计净值'].rolling(window=30).mean()
    weekly_pd_data['MA3_diff'] = weekly_pd_data['MA3'] - weekly_pd_data['MA3'].shift(1)
    pd_data = get_draw_down(pd_data)
    pd_data['MA1'] = pd_data['累计净值'].rolling(window=20).mean()
    pd_data['MA1_diff'] = pd_data['MA1'] - pd_data['MA1'].shift(1)
    pd_data['MA2'] = pd_data['累计净值'].rolling(window=30).mean()
    pd_data['MA2_diff'] = pd_data['MA2'] - pd_data['MA2'].shift(1)
    pd_data['MA3'] = pd_data['累计净值'].rolling(window=40).mean()
    pd_data['MA3_diff'] = pd_data['MA3'] - pd_data['MA3'].shift(1)
    pd_data[40:].to_excel(os.path.join(daily_cooked_path, file_name), index=False)
    weekly_pd_data[40:].to_excel(os.path.join(weekly_cooked_path, file_name), index=False)


def calculate_weekly_chat(pd_data):
    """计算周线"""
    pd.to_datetime(pd_data['净值日期'])
    pd_data.set_index('净值日期', drop=False, inplace=True)
    weekly_pd_data = pd_data.resample('W').last()
    weekly_pd_data.drop(['日增长率'], axis=1, inplace=True)
    weekly_pd_data = weekly_pd_data.dropna()
    return weekly_pd_data


def get_draw_down(pd_data):
    """最大回撤"""
    cumulative_high = 0  # 累积高点
    cumulative_high_unit = 0  # 累积高点对应的单位净值
    cumulative_high_time = None  # 累积高点对应的日期
    history_max_draw_down = 0  # 历史最大回撤
    cumulative_high_list = []
    cumulative_high_unit_list = []
    cumulative_high_time_list = []
    draw_down_list = []
    history_max_draw_down_list = []
    for _, row in pd_data.iterrows():
        if row['累计净值'] >= cumulative_high:
            cumulative_high = row['累计净值']
            cumulative_high_unit = row['单位净值']
            cumulative_high_time = row['净值日期']
        cumulative_high_list.append(cumulative_high)
        cumulative_high_unit_list.append(cumulative_high_unit)
        cumulative_high_time_list.append(cumulative_high_time)
        draw_down = round((cumulative_high - row['累计净值']) / cumulative_high_unit, 4)
        draw_down_list.append(draw_down)
        history_max_draw_down = max(history_max_draw_down, draw_down)
        history_max_draw_down_list.append(history_max_draw_down)
    pd_data['cumulative_high'] = cumulative_high_list
    pd_data['cumulative_high_unit'] = cumulative_high_unit_list
    pd_data['cumulative_high_time'] = cumulative_high_time_list
    pd_data['draw_down'] = draw_down_list
    pd_data['history_max_draw_down'] = history_max_draw_down_list

    return pd_data



# boll


# macd


# kdj